#!/usr/bin/env python

from optparse import OptionParser

def gene_list_lookup(associations_file, gene_list, options):
    if options.verbose:
        print "gene_list_lookup(%s,%s,%s)"%(associations_file, options.dataType, gene_list)

    limit = options.limit

    report = {
        "lines_processed": 0,
        "passcount": 0,
        "counts": {
            "A": 0,
            "B": 0,
            "AB": 0
        },
        "counts_by_gene": {}
    }

    for gene in gene_list: report["counts_by_gene"][gene] = 0

    rfile = open(associations_file)
    for linenumber, line in enumerate(rfile):
        colvalues = line.split("\t")

        gene1 = colvalues[0].split(":")[2]
        gene2 = colvalues[1].split(":")[2]
        spearman = colvalues[2]
        N = colvalues[3]
        log10p = colvalues[4]

        AorB = "NOT_FOUND"
        isA = (gene1 in gene_list)
        isB = (gene2 in gene_list)
        if isA and isB: AorB = "AB"
        elif isA: AorB = "A"
        elif isB: AorB = "B"

        if not AorB == "NOT_FOUND":
            print "%s\t%s\t%s\t%s\t%s\t%s" % (gene1, gene2, AorB, spearman, N, log10p)
            report["counts"][AorB] += 1
            report["passcount"] += 1

        if isA: report["counts_by_gene"][gene1] = report["counts_by_gene"][gene1] + 1
        if isB: report["counts_by_gene"][gene2] = report["counts_by_gene"][gene2] + 1

        if not limit is None and (limit <= linenumber): break
        report["lines_processed"] += 1

    generate_report(report, limit)

def generate_report(report, limit):
    report_file = open("gene_list_lookup.rpt", "w")

    report_file.write("[summary]\n")
    if not limit is None: report_file.write("limit=" + str(limit) + "\n")

    report_file.write("processed=" + str(report["lines_processed"]) + "\n")
    report_file.write("passed=" + str(report["passcount"]) + "\n")
    report_file.write("A=" + str(report["counts"]["A"]) + "\n")
    report_file.write("B=" + str(report["counts"]["B"]) + "\n")
    report_file.write("AB=" + str(report["counts"]["AB"]) + "\n")

    report_file.write("\n[counts_by_gene]\n")
    for gene in report["counts_by_gene"].keys():
        report_file.write("%s=%s\n" % (gene, report["counts_by_gene"][gene]))
    report_file.close()

if __name__ == "__main__":
    parser = OptionParser(usage="%prog ASSOCIATIONS_FILE")
    parser.add_option("-t", "--dataType", dest="dataType", default="GEXP", help="Data type for target of operation (default = GEXP)")
    parser.add_option("-g", "--geneList", dest="geneList", help="List of comma-separated genes")
    parser.add_option("-l", "--limit", dest="limit", type="int", help="Limits output, for debugging purposes")
    parser.add_option("-v", "--verbose", action="store_true", dest="verbose", default=False, help="Prints debugging statements")

    (options, args) = parser.parse_args()

    if len(args) < 1:
        parser.print_help()
        exit(-1)

    associations_file = args[0]
    if associations_file is None:
        parser.print_help()
        exit(-1)

    gene_list = options.geneList
    if gene_list is None:
        print "This program is pretty much useless without a gene list.  Use the -g option to pass in genes"
        parser.print_help()
        exit(-1)

    try:
        gene_list_lookup(associations_file, gene_list.split(","), options)
    except IOError as e:
        print e
